package com.model;

import com.persistent.PersistentChatMemoryStore;
import dev.langchain4j.http.client.HttpClientBuilder;
import dev.langchain4j.http.client.spring.restclient.SpringRestClient;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import jakarta.annotation.Resource;
import org.bsc.langgraph4j.checkpoint.BaseCheckpointSaver;
import org.bsc.langgraph4j.checkpoint.MemorySaver;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.ConfigurableApplicationContext;
import org.springframework.context.annotation.*;
import org.springframework.core.task.AsyncTaskExecutor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.web.client.RestClient;

/**
 * @Author nbZhuozh
 * @Date 2025-07-16  20:33
 **/
@Configuration
public class ModelConfig {


    @Resource
    private PersistentChatMemoryStore persistentChatMemoryStore;

    @Resource
    private ModelFactory chatModelFactory;

//    @Resource
//    private EmbeddingStoreFactory embeddingStoreFactory;

    @Resource
    private ConfigurableApplicationContext context;


//    @Lazy
//    @Bean
//    EmbeddingStore embeddingStore() {
//        return embeddingStoreFactory.getEmbeddingStore("ai_agent");
//    }

    @Lazy
    @Bean
    @Primary
    ChatModel chatLanguageModel() {
        return chatModelFactory.getChatModel("defaultModel");
    }

    @Lazy
    @Bean
    @Primary
    StreamingChatModel streamingChatLanguageModel() {
        return chatModelFactory.getStreamingChatModel("defaultModel");
    }

//    @Lazy
//    @Bean
//    @Primary
//    AiAssistant aiAssistant(){
//        return context.getBean(chatModelFactory.defaultAssistant(), AiAssistant.class);
//    }

    @Bean("languageModelHttpClientBuilder")
    @ConditionalOnMissingBean(name = "languageModelHttpClientBuilder")
    HttpClientBuilder languageModelHttpClientBuilder(ObjectProvider<RestClient.Builder> restClientBuilder) {
        return SpringRestClient.builder()
                .restClientBuilder(restClientBuilder.getIfAvailable(RestClient::builder))
                // executor is not needed for no-streaming OpenAiLanguageModel
                .createDefaultStreamingRequestExecutor(false);
    }

    @Bean("streamingLanguageModelHttpClientBuilder")
    @ConditionalOnMissingBean(name = "streamingLanguageModelHttpClientBuilder")
    HttpClientBuilder streamingLanguageModelHttpClientBuilder(
            @Qualifier("streamingLanguageModelTaskExecutor") AsyncTaskExecutor executor,
            ObjectProvider<RestClient.Builder> restClientBuilder
    ) {
        return SpringRestClient.builder()
                .restClientBuilder(restClientBuilder.getIfAvailable(RestClient::builder))
                .streamingRequestExecutor(executor);
    }

    @Bean("streamingLanguageModelTaskExecutor")
    @ConditionalOnMissingBean(name = "streamingLanguageModelTaskExecutor")
    @ConditionalOnMissingClass("io.micrometer.context.ContextSnapshotFactory")
    AsyncTaskExecutor StreamingChatModelTaskExecutor() {
        ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();
        taskExecutor.setThreadNamePrefix("languageModelTask");
        return taskExecutor;
    }

    @Bean("streamingLanguageModelTaskExecutor")
    @ConditionalOnMissingBean(name = "streamingLanguageModelTaskExecutor")
    @ConditionalOnClass(name = "io.micrometer.context.ContextSnapshotFactory")
    AsyncTaskExecutor StreamingChatModelTaskExecutorWithContextPropagation() {
        ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();
        taskExecutor.setThreadNamePrefix("languageModelTask");
        return taskExecutor;
    }

//    @Bean
//    @ConditionalOnProperty("ai.embedding-model.base-url")
//    OpenAiEmbeddingModel openAiEmbeddingModel(
//            @Qualifier("embeddingModelHttpClientBuilder") HttpClientBuilder httpClientBuilder,
//            AiProperties properties
//    ) {
//        EmbeddingModelProperties embeddingModelProperties = properties.embeddingModel();
//        return OpenAiEmbeddingModel.builder()
//                .httpClientBuilder(httpClientBuilder)
//                .baseUrl(embeddingModelProperties.baseUrl())
//                .modelName(embeddingModelProperties.modelId())
//                .apiKey(embeddingModelProperties.apiKey())
//                .dimensions(Constants.MODEL_TIME_OUT)
//                .build();
//    }

    @Bean("embeddingModelHttpClientBuilder")
    @ConditionalOnProperty("ai.embedding-model.base-url")
    @ConditionalOnMissingBean(name = "embeddingModelHttpClientBuilder")
    HttpClientBuilder openAiEmbeddingModelHttpClientBuilder(ObjectProvider<RestClient.Builder> restClientBuilder) {
        return SpringRestClient.builder()
                .restClientBuilder(restClientBuilder.getIfAvailable(RestClient::builder))
                // executor is not needed for no-streaming OpenAiEmbeddingModel
                .createDefaultStreamingRequestExecutor(false);
    }

    @Bean
    ChatMemory chatMessage(){
        return MessageWindowChatMemory
                .builder()
                .maxMessages(10)
                .chatMemoryStore(persistentChatMemoryStore)
                .build();
    }

//    @Bean
//    org.elasticsearch.client.RestClient elasticRestClient(AiProperties properties){
//        ElasticsearchProperties elasticsearchProperties = properties.elasticSearch();
//        final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
//        credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(elasticsearchProperties.userName(), elasticsearchProperties.password()));
//        return org.elasticsearch.client.RestClient.builder(HttpHost.create(elasticsearchProperties.host()))
//                .setHttpClientConfigCallback(httpClientBuilder -> {
//                    httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
//                    return httpClientBuilder;
//                })
//                .build();
//    }

    @Bean
    ChatMemoryProvider chatMemoryProvider() {
        return memoryId -> MessageWindowChatMemory
                .builder()
                .id(memoryId)
                .maxMessages(10)
                .chatMemoryStore(persistentChatMemoryStore)
                .build();
    }

    @Bean
    @Profile("dev")
    BaseCheckpointSaver memorySaver() {
        return new MemorySaver();
    }
}
